from transitions import Machine
from datetime import datetime
import logging
import json
import asyncio
import threading
from typing import Dict, List, Optional, Callable, Set
from dataclasses import dataclass, field
from enum import Enum
from queue import Queue, PriorityQueue
import time
from concurrent.futures import ThreadPoolExecutor
import uuid
import random

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

class PackageState(Enum):
    """包裹状态枚举"""
    ORDER_PLACED = "下单"
    COLLECTED = "已揽收"
    IN_TRANSIT = "运输中"
    OUT_FOR_DELIVERY = "派送中"
    DELIVERED = "已签收"
    EXCEPTION = "异常快递"

@dataclass
class PackageInfo:
    """包裹信息"""
    tracking_id: str
    sender: str
    receiver: str
    current_location: str
    created_at: datetime
    updated_at: datetime
    priority: int = 1  # 优先级：1-普通，2-加急，3-特急
    processing_time: float = 0.0  # 处理耗时统计
    state: PackageState = PackageState.ORDER_PLACED
    retry_count: int = 0
    agent_id: str = ""  # 关联的智能体ID

    def __lt__(self, other):
        """定义小于比较，用于优先队列排序"""
        if not isinstance(other, PackageInfo):
            return NotImplemented
        # 优先级高的排在前面（数字越大优先级越高，但队列是最小堆，所以用相反的比较）
        if self.priority != other.priority:
            return self.priority > other.priority
        # 优先级相同时，按创建时间排序（早创建的优先）
        return self.created_at < other.created_at

class Message:
    """智能体间通信消息"""
    def __init__(self, sender: str, receiver: str, msg_type: str, content: dict):
        self.sender = sender
        self.receiver = receiver
        self.msg_type = msg_type
        self.content = content
        self.timestamp = datetime.now()

class BaseAgent:
    """智能体基类"""
    
    def __init__(self, agent_id: str, agent_type: str):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.message_queue = Queue()
        self.is_active = False
        self.logger = logging.getLogger(f"Agent.{agent_id}")
        
    def start(self):
        """启动智能体"""
        self.is_active = True
        self.agent_thread = threading.Thread(target=self._run_loop)
        self.agent_thread.daemon = True
        self.agent_thread.start()
        self.logger.info(f"智能体 {self.agent_id} 已启动")
    
    def stop(self):
        """停止智能体"""
        self.is_active = False
        self.logger.info(f"智能体 {self.agent_id} 已停止")
    
    def _run_loop(self):
        """智能体主循环"""
        while self.is_active:
            try:
                self._process_messages()
                self._autonomous_behavior()
                time.sleep(0.1)
            except Exception as e:
                self.logger.error(f"智能体运行出错: {e}")
    
    def _process_messages(self):
        """处理消息队列"""
        try:
            while not self.message_queue.empty():
                message = self.message_queue.get_nowait()
                self.handle_message(message)
        except:
            pass
    
    def handle_message(self, message: Message):
        """处理消息 - 子类重写"""
        pass
    
    def _autonomous_behavior(self):
        """自主行为 - 子类重写"""
        pass
    
    def send_message(self, receiver: 'BaseAgent', msg_type: str, content: dict):
        """发送消息给其他智能体"""
        message = Message(self.agent_id, receiver.agent_id, msg_type, content)
        receiver.message_queue.put(message)

class PackageAgent(BaseAgent):
    """包裹智能体 - 基于流水线的智能包裹"""
    
    def __init__(self, package_info: PackageInfo, pipeline_system):
        super().__init__(f"Package_{package_info.tracking_id}", "PackageAgent")
        self.package_info = package_info
        self.package_info.agent_id = self.agent_id
        self.pipeline_system = pipeline_system
        self.current_stage = None
        self.negotiation_attempts = 0
        self.max_negotiation_attempts = 3
        
    def _autonomous_behavior(self):
        """包裹智能体的自主行为"""
        # 智能选择最优路径和处理策略
        if self.package_info.state == PackageState.ORDER_PLACED:
            self._request_collection()
        elif self.package_info.state == PackageState.EXCEPTION:
            self._handle_exception()
    
    def _request_collection(self):
        """智能请求揽收"""
        # 向最适合的揽收阶段请求处理
        best_stage = self._find_best_collection_stage()
        if best_stage:
            self.send_message(
                best_stage,
                "REQUEST_PRIORITY_PROCESSING",
                {
                    "package": self.package_info,
                    "priority": self.package_info.priority,
                    "urgency": self._calculate_urgency()
                }
            )
    
    def _find_best_collection_stage(self):
        """智能寻找最佳揽收阶段"""
        # 这里简化为返回系统的揽收阶段智能体
        return self.pipeline_system.get_stage_agent("collection")
    
    def _calculate_urgency(self) -> float:
        """计算紧急程度"""
        age = (datetime.now() - self.package_info.created_at).total_seconds()
        urgency = self.package_info.priority * 1.5 + (age / 3600)  # 基于优先级和等待时间
        return urgency
    
    def _handle_exception(self):
        """智能处理异常"""
        if self.package_info.retry_count < 3:
            # 请求重新处理
            self.send_message(
                self.pipeline_system.get_stage_agent("exception"),
                "REQUEST_RETRY",
                {
                    "package": self.package_info,
                    "retry_count": self.package_info.retry_count
                }
            )
        else:
            # 请求人工处理
            self.logger.warning(f"包裹 {self.package_info.tracking_id} 需要人工处理")
    
    def handle_message(self, message: Message):
        """处理消息"""
        if message.msg_type == "PROCESSING_ACCEPTED":
            self.current_stage = message.content["stage"]
            self.negotiation_attempts = 0
        elif message.msg_type == "PROCESSING_REJECTED":
            self.negotiation_attempts += 1
            if self.negotiation_attempts < self.max_negotiation_attempts:
                time.sleep(0.5)  # 等待后重试
        elif message.msg_type == "STATUS_UPDATE":
            self.package_info.state = PackageState(message.content["new_state"])
            self.package_info.current_location = message.content["location"]
            self.package_info.updated_at = datetime.now()

class SmartPipelineStage(BaseAgent):
    """智能流水线阶段 - 继承原有流水线功能并添加智能体能力"""
    
    def __init__(self, name: str, stage_type: str, max_workers: int = 3, processing_time: float = 1.0):
        super().__init__(f"Stage_{name}", "StageAgent")
        self.name = name
        self.stage_type = stage_type
        self.max_workers = max_workers
        self.processing_time = processing_time
        
        # 流水线功能
        self.input_queue = PriorityQueue()  # 改为优先队列
        self.output_queue = Queue()
        self.workers = []
        self.is_running = False
        self.processed_count = 0
        
        # 智能体功能
        self.load_factor = 0.0
        self.performance_history = []
        self.package_agents: Dict[str, PackageAgent] = {}
        
    def start(self):
        """启动智能流水线阶段"""
        super().start()  # 启动智能体
        
        self.is_running = True
        self.workers = []
        
        for i in range(self.max_workers):
            worker = threading.Thread(target=self._smart_worker_loop, args=(i,))
            worker.daemon = True
            worker.start()
            self.workers.append(worker)
        
        self.logger.info(f"智能阶段 {self.name} 启动，工作线程数: {self.max_workers}")
    
    def stop(self):
        """停止智能流水线阶段"""
        self.is_running = False
        super().stop()  # 停止智能体
        
        for worker in self.workers:
            worker.join(timeout=1.0)
        self.logger.info(f"智能阶段 {self.name} 已停止")
    
    def _smart_worker_loop(self, worker_id: int):
        """智能工作线程循环"""
        while self.is_running:
            try:
                # 从优先队列获取包裹
                priority, package = self.input_queue.get(timeout=0.5)
                if package is None:
                    continue
                
                start_time = time.time()
                
                # 智能处理包裹
                self.logger.info(f"智能工作线程-{worker_id} 开始处理包裹 {package.tracking_id}")
                processed_package = self._smart_process(package)
                
                # 统计处理时间
                processing_time = time.time() - start_time
                processed_package.processing_time += processing_time
                
                # 更新性能统计
                self._update_performance_stats(processing_time)
                
                # 放入输出队列
                self.output_queue.put(processed_package)
                self.processed_count += 1
                
                # 通知包裹智能体
                package_agent = self.package_agents.get(package.tracking_id)
                if package_agent:
                    self.send_message(
                        package_agent,
                        "STATUS_UPDATE",
                        {
                            "new_state": self._get_next_state().value,
                            "location": self._get_next_location(),
                            "processing_time": processing_time
                        }
                    )
                
                self.logger.info(f"智能工作线程-{worker_id} 完成处理包裹 {package.tracking_id}，耗时: {processing_time:.2f}s")
                
            except Exception as e:
                if self.is_running:
                    self.logger.error(f"智能工作线程-{worker_id} 处理出错: {e}")
    
    def _smart_process(self, package: PackageInfo) -> PackageInfo:
        """智能处理包裹"""
        # 根据包裹优先级调整处理时间
        adjusted_time = self.processing_time / package.priority
        
        # 根据当前负载调整处理策略
        if self.load_factor > 0.8:
            adjusted_time *= 0.9  # 高负载时加快处理
        
        time.sleep(adjusted_time)
        
        # 执行特定阶段处理
        if self.stage_type == "collection":
            package.current_location = "分拣中心"
            package.state = PackageState.COLLECTED
        elif self.stage_type == "transit":
            package.current_location = "目标城市分拣中心"
            package.state = PackageState.IN_TRANSIT
        elif self.stage_type == "delivery":
            if random.random() < 0.1:  # 10% 异常率
                package.current_location = "异常处理中心"
                package.state = PackageState.EXCEPTION
            else:
                package.current_location = "用户手中"
                package.state = PackageState.DELIVERED
        elif self.stage_type == "exception":
            package.current_location = "重新派送"
            package.state = PackageState.OUT_FOR_DELIVERY
            package.retry_count += 1
        
        package.updated_at = datetime.now()
        return package
    
    def _update_performance_stats(self, processing_time: float):
        """更新性能统计"""
        self.performance_history.append({
            'timestamp': datetime.now(),
            'processing_time': processing_time,
            'queue_size': self.input_queue.qsize()
        })
        
        if len(self.performance_history) > 100:
            self.performance_history.pop(0)
        
        # 更新负载因子
        self.load_factor = self.input_queue.qsize() / (self.max_workers * 2)
    
    def _get_next_state(self) -> PackageState:
        """获取下一状态"""
        state_map = {
            "collection": PackageState.COLLECTED,
            "transit": PackageState.IN_TRANSIT,
            "delivery": PackageState.DELIVERED,
            "exception": PackageState.OUT_FOR_DELIVERY
        }
        return state_map.get(self.stage_type, PackageState.ORDER_PLACED)
    
    def _get_next_location(self) -> str:
        """获取下一位置"""
        location_map = {
            "collection": "分拣中心",
            "transit": "目标城市分拣中心",
            "delivery": "用户手中",
            "exception": "重新派送"
        }
        return location_map.get(self.stage_type, "未知位置")
    
    def handle_message(self, message: Message):
        """处理智能体消息"""
        if message.msg_type == "REQUEST_PRIORITY_PROCESSING":
            package = message.content["package"]
            priority = message.content["priority"]
            urgency = message.content["urgency"]
            
            # 智能决策是否接受处理
            if self._should_accept_package(package, priority, urgency):
                # 加入优先队列
                queue_priority = -urgency  # 负数表示高优先级
                self.input_queue.put((queue_priority, package))
                
                # 记录包裹智能体
                package_agent = self.package_agents.get(package.tracking_id)
                if package_agent:
                    self.send_message(
                        package_agent,
                        "PROCESSING_ACCEPTED",
                        {"stage": self.agent_id}
                    )
                    
                self.logger.info(f"接受高优先级包裹 {package.tracking_id} 处理请求")
            else:
                # 拒绝处理
                package_agent = self.package_agents.get(package.tracking_id)
                if package_agent:
                    self.send_message(
                        package_agent,
                        "PROCESSING_REJECTED",
                        {"reason": "负载过高"}
                    )
    
    def _should_accept_package(self, package: PackageInfo, priority: int, urgency: float) -> bool:
        """智能决策是否接受包裹"""
        # 基于当前负载、包裹优先级和紧急程度决策
        if self.load_factor > 0.9 and priority == 1:
            return False  # 高负载时拒绝普通优先级
        
        if urgency > 5.0:  # 非常紧急的包裹
            return True
        
        if self.load_factor < 0.5:  # 低负载时接受所有包裹
            return True
        
        return priority >= 2  # 中等负载时只接受加急以上
    
    def add_package(self, package: PackageInfo):
        """添加包裹到队列"""
        priority = -package.priority  # 负数表示高优先级
        self.input_queue.put((priority, package))
    
    def add_package_agent(self, package_agent: PackageAgent):
        """注册包裹智能体"""
        self.package_agents[package_agent.package_info.tracking_id] = package_agent
    
    def get_processed_package(self) -> Optional[PackageInfo]:
        """获取处理完的包裹"""
        try:
            return self.output_queue.get_nowait()
        except:
            return None
    
    def get_stats(self) -> Dict:
        """获取统计信息"""
        avg_processing_time = 0
        if self.performance_history:
            avg_processing_time = sum(p['processing_time'] for p in self.performance_history) / len(self.performance_history)
        
        return {
            'name': self.name,
            'stage_type': self.stage_type,
            'processed_count': self.processed_count,
            'input_queue_size': self.input_queue.qsize(),
            'output_queue_size': self.output_queue.qsize(),
            'workers': len(self.workers),
            'is_running': self.is_running,
            'load_factor': self.load_factor,
            'avg_processing_time': avg_processing_time,
            'registered_agents': len(self.package_agents)
        }

class HybridLogisticsSystem:
    """混合智能体-流水线物流系统"""
    
    def __init__(self):
        # 创建智能流水线阶段
        self.collection_stage = SmartPipelineStage("揽收阶段", "collection", max_workers=5, processing_time=1.0)
        self.transit_stage = SmartPipelineStage("运输阶段", "transit", max_workers=10, processing_time=2.0)
        self.delivery_stage = SmartPipelineStage("派送阶段", "delivery", max_workers=8, processing_time=1.5)
        self.exception_stage = SmartPipelineStage("异常处理阶段", "exception", max_workers=3, processing_time=3.0)
        
        # 包裹和智能体管理
        self.packages: Dict[str, PackageInfo] = {}
        self.package_agents: Dict[str, PackageAgent] = {}
        self.completed_packages: List[PackageInfo] = []
        self.exception_packages: List[PackageInfo] = []
        
        # 系统控制
        self.is_running = False
        self.coordinator_thread = None
        self.logger = logging.getLogger("HybridSystem")
        
        # 性能统计
        self.total_processed = 0
        self.start_time = None
    
    def start(self):
        """启动混合系统"""
        self.logger.info("启动智能体-流水线混合物流系统...")
        
        # 启动所有智能阶段
        self.collection_stage.start()
        self.transit_stage.start()
        self.delivery_stage.start()
        self.exception_stage.start()
        
        # 启动智能协调器
        self.is_running = True
        self.start_time = time.time()
        self.coordinator_thread = threading.Thread(target=self._smart_coordinate)
        self.coordinator_thread.daemon = True
        self.coordinator_thread.start()
        
        self.logger.info("智能体-流水线混合物流系统已启动")
    
    def stop(self):
        """停止混合系统"""
        self.logger.info("停止智能体-流水线混合物流系统...")
        
        self.is_running = False
        
        # 停止所有包裹智能体
        for agent in self.package_agents.values():
            agent.stop()
        
        # 停止所有智能阶段
        self.collection_stage.stop()
        self.transit_stage.stop()
        self.delivery_stage.stop()
        self.exception_stage.stop()
        
        # 停止协调器
        if self.coordinator_thread:
            self.coordinator_thread.join(timeout=2.0)
        
        self.logger.info("智能体-流水线混合物流系统已停止")
    
    def _smart_coordinate(self):
        """智能协调器"""
        while self.is_running:
            try:
                # 智能调度包裹流转
                self._smart_package_flow()
                time.sleep(0.1)
                
            except Exception as e:
                self.logger.error(f"智能协调器出错: {e}")
    
    def _smart_package_flow(self):
        """智能包裹流转"""
        # 揽收 -> 运输
        package = self.collection_stage.get_processed_package()
        if package:
            self.transit_stage.add_package(package)
        
        # 运输 -> 派送
        package = self.transit_stage.get_processed_package()
        if package:
            self.delivery_stage.add_package(package)
        
        # 派送 -> 完成或异常
        package = self.delivery_stage.get_processed_package()
        if package:
            if package.state == PackageState.EXCEPTION:
                self.exception_packages.append(package)
                self.exception_stage.add_package(package)
            else:
                self.completed_packages.append(package)
                self.total_processed += 1
                # 停止对应的包裹智能体
                agent = self.package_agents.get(package.tracking_id)
                if agent:
                    agent.stop()
                    del self.package_agents[package.tracking_id]
                self.logger.info(f"包裹 {package.tracking_id} 完成配送，总耗时: {package.processing_time:.2f}s")
        
        # 异常处理 -> 重新派送
        package = self.exception_stage.get_processed_package()
        if package:
            self.delivery_stage.add_package(package)
    
    def add_package(self, tracking_id: str, sender: str, receiver: str, location: str, priority: int = 1):
        """添加智能包裹"""
        package = PackageInfo(
            tracking_id=tracking_id,
            sender=sender,
            receiver=receiver,
            current_location=location,
            created_at=datetime.now(),
            updated_at=datetime.now(),
            priority=priority
        )
        
        # 创建包裹智能体
        package_agent = PackageAgent(package, self)
        
        # 注册到系统
        self.packages[tracking_id] = package
        self.package_agents[tracking_id] = package_agent
        
        # 注册到各阶段
        self.collection_stage.add_package_agent(package_agent)
        self.transit_stage.add_package_agent(package_agent)
        self.delivery_stage.add_package_agent(package_agent)
        self.exception_stage.add_package_agent(package_agent)
        
        # 启动包裹智能体
        package_agent.start()
        
        # 添加到揽收阶段
        self.collection_stage.add_package(package)
        
        self.logger.info(f"新智能包裹 {tracking_id} 加入系统，优先级: {priority}")
        return package
    
    def batch_add_packages(self, packages_data: List[Dict]):
        """批量添加智能包裹"""
        for data in packages_data:
            self.add_package(**data)
    
    def get_stage_agent(self, stage_type: str):
        """获取阶段智能体"""
        stage_map = {
            "collection": self.collection_stage,
            "transit": self.transit_stage,
            "delivery": self.delivery_stage,
            "exception": self.exception_stage
        }
        return stage_map.get(stage_type)
    
    def get_system_stats(self) -> Dict:
        """获取系统统计信息"""
        runtime = time.time() - self.start_time if self.start_time else 0
        throughput = self.total_processed / (runtime / 60) if runtime > 0 else 0
        
        return {
            'runtime_seconds': runtime,
            'total_packages': len(self.packages),
            'completed_packages': len(self.completed_packages),
            'exception_packages': len(self.exception_packages),
            'active_package_agents': len(self.package_agents),
            'throughput_per_minute': throughput,
            'stages': {
                'collection': self.collection_stage.get_stats(),
                'transit': self.transit_stage.get_stats(),
                'delivery': self.delivery_stage.get_stats(),
                'exception': self.exception_stage.get_stats()
            }
        }